import numpy as np
import cpu_extension
import time

M, K, N = 1024, 1024, 1024  # 記憶體與 CPU 壓力參數，可調大測效能

A = np.random.rand(M, K).astype(np.float32)
B = np.random.rand(K, N).astype(np.float32)

start = time.time()
C = cpu_extension.matmul(A, B)
end = time.time()

print(f"CPU matrix multiplication done in {end - start:.3f} sec")
print(f"Result shape: {C.shape}")